from autogen import AssistantAgent, UserProxyAgent
import requests
import json

class DouBaoLLM:
    def __init__(self,api_key):
        self.api_key = '9d4b99d0-d0a4-4212-9e4d-0ed00c0d313a'
        self.model = 'deepseek-v3-1-terminus'
        self.base_url = "https://ark.cn-beijing.volces.com/api/v3"

    def create_chat_completion(self, messages, **kwargs):
        """模拟 OpenAI 格式的聊天完成接口"""
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }

        payload = {
            "model": self.model,
            "messages": messages,
            "temperature": kwargs.get("temperature", 0.7),
            "stream": False
        }

        try:
            response = requests.post(
                f"{self.base_url}/chat/completions",
                headers=headers,
                json=payload,
                timeout=30
            )
            response.raise_for_status()
            result = response.json()

            # 返回与 OpenAI 兼容的格式
            return {
                "choices": [{
                    "message": result["choices"][0]["message"],
                    "finish_reason": "stop"
                }]
            }
        except Exception as e:
            print(f"豆包API调用失败: {e}")
            return {"choices": [{"message": {"content": "抱歉，我暂时无法回应。"}}]}


# 创建豆包增强的 AutoGen 智能体
def create_doubao_agents(api_key: str):
    """创建使用豆包模型的 AutoGen 智能体"""

    # 豆包配置
    doubao_config = {
        "config_list": [{
            "model": "deepseek-v3-1-terminus",
            "api_key": api_key,
            "client": DouBaoLLM(api_key)
        }],
        "temperature": 0.7,
        "timeout": 30
    }

    # 创建助理智能体
    assistant = AssistantAgent(
        name="doubao_assistant",
        system_message="你是一个有帮助的AI助手，基于豆包大模型。",
        llm_config=doubao_config,
    )

    # 创建用户代理
    user_proxy = UserProxyAgent(
        name="user_proxy",
        human_input_mode="NEVER",
        max_consecutive_auto_reply=5,
        code_execution_config={"work_dir": "coding"}
    )

    return user_proxy, assistant


# 使用示例
def run_doubao_autogen():
    api_key = "9d4b99d0-d0a4-4212-9e4d-0ed00c0d313a"  # 替换为您的豆包 API Key

    user_proxy, assistant = create_doubao_agents(api_key)

    # 开始对话
    user_proxy.initiate_chat(
        assistant,
        message="请用Python写一个快速排序算法，并解释其工作原理。"
    )

# 运行示例
run_doubao_autogen()